# Author: Cascade
# Date: 2024/06/24
# Description: CosyVoice V2 TTS Service Implementation using Alibaba Cloud DashScope API

import os
import logging
import time
import base64
from typing import List, Dict, Any, Optional
from pathlib import Path

import dashscope
from dashscope.audio.tts_v2 import SpeechSynthesizer, AudioFormat

from pptflow.tts.tts_service import TtsService
from pptflow.config.setting import Setting
from pptflow.utils import mylogger

logger = mylogger.get_logger(__name__)

class CosyVoiceV2TtsService(TtsService):
    """
    CosyVoice V2 TTS Service Implementation using Alibaba Cloud DashScope API
    """
    
    # Audio format mapping from string to AudioFormat enum
    AUDIO_FORMATS = {
        'wav_8k': AudioFormat.WAV_8000HZ_MONO_16BIT,
        'wav_16k': AudioFormat.WAV_16000HZ_MONO_16BIT,
        'wav_22k': AudioFormat.WAV_22050HZ_MONO_16BIT,
        'wav_24k': AudioFormat.WAV_24000HZ_MONO_16BIT,
        'wav_44k': AudioFormat.WAV_44100HZ_MONO_16BIT,
        'wav_48k': AudioFormat.WAV_48000HZ_MONO_16BIT,
        'mp3_8k': AudioFormat.MP3_8000HZ_MONO_128KBPS,
        'mp3_16k': AudioFormat.MP3_16000HZ_MONO_128KBPS,
        'mp3_22k': AudioFormat.MP3_22050HZ_MONO_256KBPS,
        'mp3_24k': AudioFormat.MP3_24000HZ_MONO_256KBPS,
        'mp3_44k': AudioFormat.MP3_44100HZ_MONO_256KBPS,
        'mp3_48k': AudioFormat.MP3_48000HZ_MONO_256KBPS,
        'pcm_8k': AudioFormat.PCM_8000HZ_MONO_16BIT,
        'pcm_16k': AudioFormat.PCM_16000HZ_MONO_16BIT,
        'pcm_22k': AudioFormat.PCM_22050HZ_MONO_16BIT,
        'pcm_24k': AudioFormat.PCM_24000HZ_MONO_16BIT,
        'pcm_44k': AudioFormat.PCM_44100HZ_MONO_16BIT,
        'pcm_48k': AudioFormat.PCM_48000HZ_MONO_16BIT,
    }
    
    def __init__(self):
        # Initialize DashScope API key
        self.api_key = os.getenv("COSYVOICE_API_KEY")
        if not self.api_key:
            raise ValueError("COSYVOICE_API_KEY environment variable is not set")
        
        # Set the API key for DashScope
        dashscope.api_key = self.api_key
        
        # Voice settings
        self.model = os.getenv("COSYVOICE_MODEL", "cosyvoice-v2")
        self.default_voice = os.getenv("COSYVOICE_DEFAULT_VOICE", "longxiaochun_v2")
        self.default_speed = float(os.getenv("COSYVOICE_SPEED", "1.0"))
        
        # Audio format settings
        audio_format = os.getenv("COSYVOICE_AUDIO_FORMAT", "wav_24k").lower()
        self.audio_format = self.AUDIO_FORMATS.get(
            audio_format, 
            AudioFormat.WAV_24000HZ_MONO_16BIT
        )
        
        # Other settings
        self.max_retries = int(os.getenv("COSYVOICE_MAX_RETRIES", "3"))
        self.retry_delay = float(os.getenv("COSYVOICE_RETRY_DELAY", "1.0"))
        
        logger.info(f"Initialized CosyVoice V2 TTS Service with model: {self.model}, voice: {self.default_voice}")
    
    async def tts(self, text: str, output_audio_filename: str, setting: Setting):
        """
        Convert text to speech using CosyVoice V2
        
        Args:
            text: The text to convert to speech
            output_audio_filename: Path to save the generated audio
            setting: Application settings
        """
        try:
            if not text.strip():
                logger.warning("Empty text provided for TTS")
                return
                
            # Create output directory if it doesn't exist
            output_path = Path(output_audio_filename)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            
            # Initialize the synthesizer
            synthesizer = SpeechSynthesizer(
                model=self.model,
                voice=self.default_voice,
                format=self.audio_format,
            )
            
            # Make the API call with retry logic
            for attempt in range(self.max_retries + 1):
                try:
                    # Call the TTS API
                    audio_data = synthesizer.call(text)
                    
                    # Save the audio file
                    with open(output_path, 'wb') as f:
                        f.write(audio_data)
                    
                    logger.info(f"TTS audio saved to {output_audio_filename}")
                    logger.debug(f"Request ID: {synthesizer.get_last_request_id()}, "
                                 f"First package delay: {synthesizer.get_first_package_delay()}ms")
                    
                    return  # Success, exit the retry loop
                    
                except Exception as e:
                    if attempt == self.max_retries:
                        logger.error(f"Failed to generate TTS after {self.max_retries} attempts: {str(e)}")
                        raise
                    
                    # Wait before retrying
                    wait_time = self.retry_delay * (2 ** attempt)  # Exponential backoff
                    logger.warning(f"Attempt {attempt + 1} failed. Retrying in {wait_time:.1f} seconds...")
                    time.sleep(wait_time)
        
        except Exception as e:
            logger.error(f"Error in CosyVoice V2 TTS: {str(e)}")
            raise
    
    def get_voice_list(self, setting: Setting = None, force_refresh: bool = False) -> List[Dict[str, Any]]:
        """
        Get the list of available voices from CosyVoice V2
        
        Args:
            setting: Application settings
            force_refresh: If True, force refresh the voice list
            
        Returns:
            List of available voices with their properties
        """
        # Note: The DashScope SDK doesn't provide an API to list available voices,
        # so we return a static list of known voices
        return [
            {
                'id': 'longxiaochun_v2',
                'name': 'LongXiaoChun V2',
                'gender': 'Female',
                'language': 'zh-CN',
                'language_name': 'Chinese (Mandarin)',
                'provider': 'cosyvoice_v2',
                'neural': True,
                'description': 'LongXiaoChun V2 voice - Chinese (Mandarin)'
            },
            {
                'id': 'longxiaochun_v1',
                'name': 'LongXiaoChun V1',
                'gender': 'Female',
                'language': 'zh-CN',
                'language_name': 'Chinese (Mandarin)',
                'provider': 'cosyvoice_v2',
                'neural': True,
                'description': 'LongXiaoChun V1 voice - Chinese (Mandarin)'
            },
            {
                'id': 'zhimiao_v2',
                'name': 'ZhiMiao V2',
                'gender': 'Female',
                'language': 'zh-CN',
                'language_name': 'Chinese (Mandarin)',
                'provider': 'cosyvoice_v2',
                'neural': True,
                'description': 'ZhiMiao V2 voice - Chinese (Mandarin)'
            },
            {
                'id': 'zhimiao_v1',
                'name': 'ZhiMiao V1',
                'gender': 'Female',
                'language': 'zh-CN',
                'language_name': 'Chinese (Mandarin)',
                'provider': 'cosyvoice_v2',
                'neural': True,
                'description': 'ZhiMiao V1 voice - Chinese (Mandarin)'
            },
            # Add more voices as needed
        ]

# For testing
if __name__ == "__main__":
    import asyncio
    from dotenv import load_dotenv
    
    async def test():
        load_dotenv()
        
        # Initialize the service
        service = CosyVoiceV2TtsService()
        
        # Test TTS
        output_file = "test_cosyvoice_v2.mp3"
        await service.tts("这是一个测试，测试CosyVoice V2语音合成效果。", output_file, None)
        print(f"TTS output saved to {output_file}")
        
        # Test voice list
        voices = service.get_voice_list()
        print("\nAvailable voices:")
        for voice in voices:
            print(f"- {voice['name']} ({voice['id']}): {voice['description']}")
    
    asyncio.run(test())
